#!/usr/bin/env python3
"""
Hadoop 3.1 MCP Server

面向 Hadoop 3.1 分布式计算框架的命令行操作 MCP 服务器。

MCP tools 列表：
1. hadoop_service_control - 控制 Hadoop 各种服务的启动、停止、重启和重载
   参数：service_name (str), action (str)
   返回：统一的命令执行结果 JSON

2. hadoop_zkfc_format - 格式化 ZooKeeper Failover Controller
   参数：无
   返回：统一的命令执行结果 JSON
"""

from mcp.server.fastmcp import FastMCP, Context
import subprocess
import shlex

# 创建 MCP 服务器实例
mcp = FastMCP("Hadoop MCP Server")

@mcp.tool()
def hadoop_service_control(service_name: str, action: str) -> dict:
    """
    控制 Hadoop 服务的启动、停止、重启和重载操作
    
    参数:
        service_name: Hadoop 服务名称，如 hadoop-datanode.service, hadoop-namenode.service 等
        action: 操作类型，支持 'start', 'stop', 'restart', 'reload'
    
    返回:
        JSON 结构:
        {
            "success": bool,      # 命令是否成功执行
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出内容
            "stderr": str         # 标准错误内容
        }
    """
    valid_actions = ['start', 'stop', 'restart', 'reload']
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} {service_name}",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"无效的操作类型: {action}，支持的操作: {', '.join(valid_actions)}"
        }
    
    cmd = f"systemctl {action} {service_name}"
    try:
        result = subprocess.run(
            shlex.split(cmd),
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": "命令执行超时"
        }
    except Exception as e:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"命令执行异常: {str(e)}"
        }

@mcp.tool()
def hadoop_zkfc_format() -> dict:
    """
    格式化 ZooKeeper Failover Controller
    
    参数:
        无
    
    返回:
        JSON 结构:
        {
            "success": bool,      # 命令是否成功执行
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出内容
            "stderr": str         # 标准错误内容
        }
    """
    cmd = "echo 'Y' | hdfs zkfc -formatZK"
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=60
        )
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": "命令执行超时"
        }
    except Exception as e:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"命令执行异常: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()